package com.fengwk.deeplearning.test;

import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import javax.imageio.ImageIO;

import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

import com.fengwk.deeplearning.core.Layer;
import com.fengwk.deeplearning.core.NeuralNetwork;
import com.fengwk.deeplearning.core.Samples;
import com.fengwk.deeplearning.core.UnitCompute;
import com.fengwk.support.util.ArrayUtils;
import com.fengwk.support.util.FileExtensionUtils;
import com.fengwk.support.util.FileUtils;
import com.fengwk.support.util.IoUtils;
import com.fengwk.support.util.LoggerUtils;
import com.fengwk.support.util.LoggerUtils.Logger;

public class CatVsDog {

	private static final Logger LOG = LoggerUtils.getLogger(CatVsDog.class);

	public static final double Y_DOG = 0;
	public static final double Y_CAT = 1;
	
	public static void resizeImg(String dirPath) throws IOException {
		File dir = FileUtils.buildFile(dirPath);
		File[] fs = dir.listFiles();
		int avgW = 0, avgH = 0;
		int count = 0;
		for (File f: fs) {
			BufferedImage bi = ImageIO.read(f);
			avgW += bi.getWidth();
			avgH += bi.getHeight();
			LOG.debug("已读取第" + count ++ + "张图片");
		}
		avgW /= count;
		avgH /= count;
		count = 0;
		for (File f: fs) {
			BufferedImage bi = ImageIO.read(f);
			BufferedImage buf = new BufferedImage(avgW, avgH, bi.getType());
			Graphics2D g = buf.createGraphics();
			g.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR);
			g.drawImage(bi, 0, 0, avgW, avgH, null);
			g.dispose();
			ImageIO.write(buf, FileExtensionUtils.split(f.getName())[1], f);
			LOG.debug("处理取第" + count ++ + "张图片");
		}
	}
	
	public static void readData(String dirPath, String dataOutPath) throws IOException {
		File dir = FileUtils.buildFile(dirPath);
		File[] fs = dir.listFiles();
		int n = 0;
		int m = 0;
		List<double[]> xs = new ArrayList<>();
		double[] ys = new double[fs.length];
		for (File f: fs) {
			BufferedImage bi = ImageIO.read(f);
			if (n == 0) {
				int w = bi.getWidth();
				int h = bi.getHeight();
				n = w * h * 3;
			}
			xs.add(readRgb(bi));
			String fileName = f.getName();
			if (fileName.contains("cat"))
				ys[m ++] = Y_CAT;
			else
				ys[m ++] = Y_DOG;
			LOG.debug("处理取第" + m + "张图片");
		}
		double[] xsArr = ArrayUtils.mergeDoubleArryList(xs);
		xs = null;
		INDArray XT = Nd4j.create(xsArr, new int[] {m, n});
		INDArray X = XT.transpose();
		XT = null;
		INDArray y = Nd4j.create(ys);
		Samples samples = new Samples(X, y);
		samples.serialize(dataOutPath);
	}
	
	private static double[] readRgb(BufferedImage bi) {
		int w = bi.getWidth();
		int h = bi.getHeight();
		int i = 0;
		double[] r0s = new double[w * h];
		double[] g0s = new double[w * h];
		double[] b0s = new double[w * h];
		for (int r = 0; r < h; r ++) {
			for (int c = 0; c < h; c ++) {
				int rgb = bi.getRGB(r, c);
				double r0 = (rgb & 0xFF0000) >> 16;
				double g0 = (rgb & 0xFF00) >> 8;
				double b0 = rgb & 0xFF;
				r0s[i] = r0;
				g0s[i] = g0;
				b0s[i ++] = b0;
			}
		}
		double[] x = ArrayUtils.merge(r0s, g0s, b0s);
		r0s = null;
		g0s = null;
		b0s = null;
		return x;
	}
	
	public static void training(String trainDataPath, String savePath) {
		Samples samples = (Samples) IoUtils.bytesToObject(IoUtils.readBytes(trainDataPath));
//		Layer hidden1 = new Layer(UnitCompute.Enum.RELU, 10);
//		Layer hidden2 = new Layer(UnitCompute.Enum.RELU, 5);
//		Layer hidden3 = new Layer(UnitCompute.Enum.SIGMOID, 1);
		Layer output = new Layer(UnitCompute.Enum.SIGMOID, 1);
		NeuralNetwork nn = new NeuralNetwork(
				null, output, 0.001f, 0f, 1f);
		nn.training(samples, 2000);
		nn.serialize(savePath);
	}
	
	public static void run(String networkPath, String path) throws IOException {
		NeuralNetwork nn = (NeuralNetwork) IoUtils.bytesToObject(IoUtils.readBytes(networkPath));
		BufferedImage bi = ImageIO.read(FileUtils.buildFile(path));
		double[] x = readRgb(bi);
		INDArray X = Nd4j.create(x, new int[] {x.length, 1});
		System.out.println(nn.run(X));
	}
	
	public static void main(String[] args) throws IOException {
//		resizeImg("d:/dataset/dataset_kaggledogvscat/test");
		readData("d:/dataset/dataset_kaggledogvscat/train-small", "d:/dataset/dataset_kaggledogvscat/train-small/data");
//		training("d:/dataset/dataset_kaggledogvscat/train-small/data", "d:/dataset/dataset_kaggledogvscat/train-small/network");
//		run("d:/dataset/dataset_kaggledogvscat/train-small/network", "d:/dataset/dataset_kaggledogvscat/test/dog.7456.jpg");
	}
	
}
